"""
可行性研究 v2：基于重构后的模块实现跳层推理实验
增强版：支持单层跳跃和连续多层跳跃实验

使用方式：
1. 直接运行：使用代码中的默认配置
2. 使用预定义配置：导入experiment_configs并指定配置名称
"""
from llm import ModelLoader, SkipableModel

# 可选：导入预定义配置 (取消注释以使用)
# from experiment_configs import get_config


def validate_consecutive_skip_config(layer_range, k_value, num_layers):
    """
    验证连续跳层实验配置的有效性
    
    Args:
        layer_range: 层数闭区间 [start, end]
        k_value: 连续跳过的层数
        num_layers: 模型总层数
        
    Returns:
        list: 所有有效的跳层范围组合 [(start1, end1), (start2, end2), ...]
        
    Raises:
        ValueError: 配置无效时抛出异常
    """
    start_range, end_range = layer_range
    
    # 参数验证
    if start_range < 0 or end_range >= num_layers:
        raise ValueError(f"层范围[{start_range}, {end_range}]超出模型层数范围[0, {num_layers-1}]")
    
    if start_range > end_range:
        raise ValueError(f"起始层({start_range})不能大于结束层({end_range})")
    
    if k_value <= 0:
        raise ValueError(f"k值({k_value})必须大于0")
    
    if k_value > (end_range - start_range + 1):
        raise ValueError(f"k值({k_value})不能大于指定范围的层数({end_range - start_range + 1})")
    
    # 生成所有有效的连续跳层组合
    consecutive_ranges = []
    for start in range(start_range, end_range - k_value + 2):
        end = start + k_value - 1
        if end <= end_range:
            consecutive_ranges.append((start, end))
    
    if not consecutive_ranges:
        raise ValueError(f"无法在范围[{start_range}, {end_range}]中生成k={k_value}的连续跳层组合")
    
    return consecutive_ranges


def main():
    """
    执行跳层推理实验 - 基于重构后的架构
    支持单层跳跃和连续多层跳跃实验
    """
    # ============================================================================
    # 实验配置区 - 可根据需要修改以下参数
    # ============================================================================
    
    # 选择配置方式：
    # 方式1：使用预定义配置 (取消注释以下两行)
    # config = get_config("comprehensive")  # 可选: "small_range_test", "middle_range_test", "large_range_test", "consecutive_only", "single_only", "comprehensive"
    # LAYER_RANGE, K_VALUE, RUN_SINGLE_LAYER_SKIP, RUN_CONSECUTIVE_SKIP, VERBOSE = config["LAYER_RANGE"], config["K_VALUE"], config["RUN_SINGLE_LAYER_SKIP"], config["RUN_CONSECUTIVE_SKIP"], config["VERBOSE"]
    
    # 方式2：直接在这里配置 (默认方式)
    LAYER_RANGE = [1, 16]  # 层数闭区间 [start, end] - 指定要测试的层范围
    K_VALUE = 3  # 连续跳过的层数
    RUN_SINGLE_LAYER_SKIP = False   # 是否执行单层跳跃实验
    RUN_CONSECUTIVE_SKIP = True    # 是否执行连续跳层实验
    VERBOSE = True  # 是否显示详细输出
    
    # 模型配置
    MODEL_ID = "meta-llama/Llama-2-13b-chat-hf"
    CACHE_DIR = "./hf_cache/"
    
    # ============================================================================
    # 实验配置区结束
    # ============================================================================
    
    # --- 1. 加载模型 ---
    print("正在加载模型...")
    loader = ModelLoader(cache_dir=CACHE_DIR)
    model, tokenizer = loader.load_model(MODEL_ID)
    
    # 创建可跳层模型
    skipable_model = SkipableModel(model, tokenizer)
    
    print(f"模型总层数: {skipable_model.num_layers}")
    if RUN_CONSECUTIVE_SKIP:
        print(f"连续跳层实验配置: 区间[{LAYER_RANGE[0]}, {LAYER_RANGE[1]}], k={K_VALUE}")
        
        # 验证配置并生成跳层组合
        try:
            consecutive_ranges = validate_consecutive_skip_config(LAYER_RANGE, K_VALUE, skipable_model.num_layers)
            print(f"将测试的跳层范围: {[f'[{start},{end}]' for start, end in consecutive_ranges]}")
            print(f"总共 {len(consecutive_ranges)} 种跳层组合")
        except ValueError as e:
            print(f"配置错误: {e}")
            return
    else:
        consecutive_ranges = []
    
    # --- 2. 准备测试用例 ---
    test_prompts = [
        # 逻辑推理问题
        "Which name is longer? Bob has 3 letters, Alice has 5 letters. Therefore, Alice is",
        
        # 数学比较问题  
        "Compare 8.11 and 8.9: Since 8.11 has more digits after decimal but 8.9 = 8.90, we see 8.9 is actually",
        
        # 算术问题
        "What is 15 + 27? Let me calculate step by step: 15 + 27 = 10 + 5 + 20 + 7 = 30 + 12 = ",
        "print('Hello ",
        "int a = 1.0f"
    ]
    
    # --- 3. 执行实验 ---
    num_layers = skipable_model.num_layers
    all_results = []
    
    for prompt_idx, prompt in enumerate(test_prompts):
        print(f"\n" + "="*70)
        print(f" 测试例子 {prompt_idx + 1}: '{prompt}' ")
        print("="*70)
        
        # 完整推理 (基准)
        print("\n执行完整推理 (基准)")
        print("-" * 30)
        
        baseline_token = skipable_model.full_inference(prompt)
        if VERBOSE:
            print(f"提示: '{prompt}'")
        print(f"基准模型预测的下一个词元: '{baseline_token}'")
        
        # 记录当前prompt的结果
        prompt_results = {
            'prompt': prompt,
            'baseline_token': baseline_token,
            'single_skip_results': [] if RUN_SINGLE_LAYER_SKIP else None,
            'consecutive_skip_results': [] if RUN_CONSECUTIVE_SKIP else None,
            'single_error_layers': [] if RUN_SINGLE_LAYER_SKIP else None,
            'consecutive_error_ranges': [] if RUN_CONSECUTIVE_SKIP else None
        }
        
        # 1. 单层跳过实验
        if RUN_SINGLE_LAYER_SKIP:
            print("\n执行单层跳层推理实验")
            print("-" * 30)
            
            for layer_to_skip in range(num_layers):
                skipped_token = skipable_model.skip_single_layer(prompt, layer_to_skip)
                is_same = (skipped_token == baseline_token)
                
                # 记录结果
                prompt_results['single_skip_results'].append({
                    'layer_skipped': layer_to_skip,
                    'predicted_token': skipped_token,
                    'is_correct': is_same
                })
                
                if not is_same:
                    prompt_results['single_error_layers'].append(layer_to_skip)
                
                if VERBOSE:
                    print(f"跳过第 {layer_to_skip:02d} 层 | 预测词元: '{skipped_token}' | "
                          f"与基准一致: {'✅' if is_same else '❌'}")
                elif not is_same:  # 只显示错误的情况
                    print(f"跳过第 {layer_to_skip:02d} 层 | 预测词元: '{skipped_token}' | "
                          f"与基准一致: {'❌'}")
        
        # 2. 连续多层跳过实验
        if RUN_CONSECUTIVE_SKIP:
            print(f"\n执行连续跳层推理实验 (k={K_VALUE})")
            print("-" * 30)
            
            for start, end in consecutive_ranges:
                skipped_token = skipable_model.skip_consecutive_layers(prompt, start, end)
                is_same = (skipped_token == baseline_token)
                
                # 记录结果
                prompt_results['consecutive_skip_results'].append({
                    'range_skipped': (start, end),
                    'predicted_token': skipped_token,
                    'is_correct': is_same
                })
                
                if not is_same:
                    prompt_results['consecutive_error_ranges'].append((start, end))
                
                if VERBOSE:
                    print(f"跳过第 [{start:02d}-{end:02d}] 层 | 预测词元: '{skipped_token}' | "
                          f"与基准一致: {'✅' if is_same else '❌'}")
                elif not is_same:  # 只显示错误的情况
                    print(f"跳过第 [{start:02d}-{end:02d}] 层 | 预测词元: '{skipped_token}' | "
                          f"与基准一致: {'❌'}")
        
        all_results.append(prompt_results)
    
    # --- 4. 统计分析 ---
    print("\n" + "="*80)
    print(" 综合统计分析 ")
    print("="*80)
    
    # 初始化统计变量
    total_single_tests = len(test_prompts) * num_layers if RUN_SINGLE_LAYER_SKIP else 0
    total_single_errors = 0
    all_single_error_layers = []
    
    total_consecutive_tests = len(test_prompts) * len(consecutive_ranges) if RUN_CONSECUTIVE_SKIP else 0
    total_consecutive_errors = 0
    all_consecutive_error_ranges = []
    
    for result in all_results:
        # 单层统计
        if RUN_SINGLE_LAYER_SKIP and result['single_error_layers'] is not None:
            single_error_count = len(result['single_error_layers'])
            total_single_errors += single_error_count
            all_single_error_layers.extend(result['single_error_layers'])
        else:
            single_error_count = 0
        
        # 连续跳层统计
        if RUN_CONSECUTIVE_SKIP and result['consecutive_error_ranges'] is not None:
            consecutive_error_count = len(result['consecutive_error_ranges'])
            total_consecutive_errors += consecutive_error_count
            all_consecutive_error_ranges.extend(result['consecutive_error_ranges'])
        else:
            consecutive_error_count = 0
        
        print(f"\n提示: '{result['prompt'][:50]}{'...' if len(result['prompt']) > 50 else ''}'")
        print(f"  基准预测: '{result['baseline_token']}'")
        
        if RUN_SINGLE_LAYER_SKIP:
            print(f"  单层跳跃错误: {single_error_count}/{num_layers} ({single_error_count/num_layers*100:.1f}%)")
            if result['single_error_layers']:
                print(f"  单层错误层序号: {result['single_error_layers']}")
        
        if RUN_CONSECUTIVE_SKIP:
            print(f"  连续跳层错误: {consecutive_error_count}/{len(consecutive_ranges)} ({consecutive_error_count/len(consecutive_ranges)*100:.1f}%)")
            if result['consecutive_error_ranges']:
                print(f"  连续跳层错误范围: {result['consecutive_error_ranges']}")
    
    # 单层跳跃统计
    if RUN_SINGLE_LAYER_SKIP:
        single_error_rate = total_single_errors / total_single_tests * 100 if total_single_tests > 0 else 0
        
        print(f"\n" + "="*60)
        print(" 单层跳跃统计 ")
        print("="*60)
        print(f"总测试次数: {total_single_tests}")
        print(f"总错误次数: {total_single_errors}")
        print(f"总体错误率: {single_error_rate:.2f}%")
        
        if all_single_error_layers:
            mean_error_layer = sum(all_single_error_layers) / len(all_single_error_layers)
            print(f"错误层序号均值: {mean_error_layer:.2f}")
            
            # 计算错误层的分布
            layer_error_count = {}
            for layer in all_single_error_layers:
                layer_error_count[layer] = layer_error_count.get(layer, 0) + 1
            
            print(f"\n单层错误分布:")
            for layer in sorted(layer_error_count.keys()):
                count = layer_error_count[layer]
                percentage = count / len(all_single_error_layers) * 100
                print(f"  第{layer:02d}层: {count}次 ({percentage:.1f}%)")
        else:
            print("错误层序号均值: 无错误")
            print("所有单层跳跃测试都与基准预测一致！")
    
    # 连续跳层统计
    if RUN_CONSECUTIVE_SKIP:
        consecutive_error_rate = total_consecutive_errors / total_consecutive_tests * 100 if total_consecutive_tests > 0 else 0
        
        print(f"\n" + "="*60)
        print(f" 连续跳层统计 (k={K_VALUE}) ")
        print("="*60)
        print(f"总测试次数: {total_consecutive_tests}")
        print(f"总错误次数: {total_consecutive_errors}")
        print(f"总体错误率: {consecutive_error_rate:.2f}%")
        
        if all_consecutive_error_ranges:
            # 计算错误范围的分布
            range_error_count = {}
            for start, end in all_consecutive_error_ranges:
                range_key = f"[{start}-{end}]"
                range_error_count[range_key] = range_error_count.get(range_key, 0) + 1
            
            print(f"\n连续跳层错误分布:")
            for range_key in sorted(range_error_count.keys()):
                count = range_error_count[range_key]
                percentage = count / len(all_consecutive_error_ranges) * 100
                print(f"  {range_key}: {count}次 ({percentage:.1f}%)")
            
            # 分析哪些层在连续跳跃中最容易出错
            involved_layers = []
            for start, end in all_consecutive_error_ranges:
                involved_layers.extend(range(start, end + 1))
            
            if involved_layers:
                layer_involvement_count = {}
                for layer in involved_layers:
                    layer_involvement_count[layer] = layer_involvement_count.get(layer, 0) + 1
                
                print(f"\n连续跳层中涉及的错误层频次:")
                for layer in sorted(layer_involvement_count.keys()):
                    count = layer_involvement_count[layer]
                    percentage = count / len(involved_layers) * 100
                    print(f"  第{layer:02d}层: {count}次 ({percentage:.1f}%)")
        else:
            print("所有连续跳层测试都与基准预测一致！")
    
    # 实验总结
    print(f"\n" + "="*60)
    print(" 实验总结 ")
    print("="*60)
    if RUN_SINGLE_LAYER_SKIP and RUN_CONSECUTIVE_SKIP:
        print(f"单层跳跃错误率: {single_error_rate:.2f}%")
        print(f"连续跳层错误率: {consecutive_error_rate:.2f}%")
        print(f"连续跳层相对于单层跳跃的错误率变化: {consecutive_error_rate - single_error_rate:+.2f}%")
    elif RUN_SINGLE_LAYER_SKIP:
        print(f"仅执行单层跳跃实验，错误率: {single_error_rate:.2f}%")
    elif RUN_CONSECUTIVE_SKIP:
        print(f"仅执行连续跳层实验，错误率: {consecutive_error_rate:.2f}%")
    
    print(f"测试的提示词数量: {len(test_prompts)}")
    print(f"实验完成！")


if __name__ == "__main__":
    main()
